import torch
from torch_geometric.data import InMemoryDataset
from torch_geometric.data import Data
from torch_geometric.data import DataLoader
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
import torch.nn as nn

class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['EDGE_FEATURES.txt', 'EDGE_INDEX.txt', 'GRAPH_INDICATOR.txt','NODE_FEATURES.txt']

    @property
    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        data_list=[]
        for i in range(0,15):
            edge_index = torch.tensor([[0, 1, 1, 2],
                                       [1, 0, 2, 1]], dtype=torch.long)
            x = torch.tensor([[i], [i+1], [i+2]], dtype=torch.float)
            y=torch.tensor([-1], dtype=torch.long)
            if i%2==0:
                y=torch.tensor([1,1,1], dtype=torch.long)
            else:
                y = torch.tensor([2,2,2], dtype=torch.long)
            data = Data(x=x,y=y, edge_index=edge_index)
            data_list.append(data)




        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])





class GCNNet(torch.nn.Module):
    # 网络初始化
    def __init__(self, input_feature=4, num_class=1):
        '''
        :param input_feature: 每个节点属性的维度
        :param num_class: 节点所属的类别数
        '''
        super(GCNNet, self).__init__()
        self.input_feature = input_feature
        self.num_classes = num_class
        # 图卷积层
        # 出自于论文《Semi-Supervised Classification with Graph Convolutional Networks》
        # 具体计算公式可查阅论文
        self.conv1 = GCNConv(input_feature, 32)
        self.conv2 = GCNConv(32, num_class)

    # 前向传递过程
    def forward(self, data):
        # data是输入到网络中的训练数据
        # 获取节点的属性信息和边的连接信息
        x, edge_index = data.x, data.edge_index
        # 卷积，利用相邻节点的信息更新自己的特征表示
        x = self.conv1(x, edge_index)
        # 激活
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        # 预测类别
        output = F.softmax(x, dim=1)
        return output






if __name__ == '__main__':
    dataset= MyOwnDataset(root='my_dataset')
    print(dataset)
    print(dataset.num_classes)
    print(dataset.num_node_features)
    print(dataset[0])
    print(dataset[0].is_undirected())
    # for data in dataset:
    #     for i in data:
    #         print(i)
    # dataloader=Dataloader(dataset=dataset)
    # print(dataloader)

    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GCNNet()
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    train_loader = DataLoader(dataset=dataset)
    # Train the model
    total_step = len(train_loader)
    num_epochs=200
    for epoch in range(num_epochs):
        for i,batch in enumerate(train_loader):
            print(batch)
            images = images.to(device)
            labels = labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimizer
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 100 == 0:
                print('Epoch: [{}/{}], Step: [{}/{}], Loss: {}'
                      .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))
